#pragma once
#include "Common.hpp"
#include "ObjectPool.hpp"

// Single-level array
// 单层基数树，按页号直接哈希
// 大小：32位下：4*2^20=4MB；64位下：4*2^52就太大了
template <int BITS> // 存储页号需要多少位 32位下：32-PAGE_SHIFT=20 64位下：64-PAGE_SHIFT=52
class TCMalloc_PageMap1
{
private:
	static const int LENGTH = 1 << BITS; // 数组长度（元素个数）
	void **array_;

public:
	typedef uintptr_t Number;

	// explicit TCMalloc_PageMap1(void* (*allocator)(size_t)) {
	explicit TCMalloc_PageMap1()
	{
		// array_ = reinterpret_cast<void**>((*allocator)(sizeof(void*) << BITS));
		size_t size = sizeof(void *) << BITS;						   // 数组大小
		size_t alignSize = AlineRule::_RoundUp(size, 1 << PAGE_SHIFT); // 按页对齐
		array_ = (void **)SystemAlloc(alignSize >> PAGE_SHIFT);		   // 申请
		memset(array_, 0, sizeof(void *) << BITS);					   // 清空
	}

	// Return the current value for KEY.  Returns NULL if not yet set,
	// or if k is out of range.
	void *get(Number k) const
	{
		if ((k >> BITS) > 0)
		{
			return NULL;
		}
		return array_[k]; // 直接哈希
	}

	// REQUIRES "k" is in range "[0,2^BITS-1]".
	// REQUIRES "k" has been ensured before.
	//
	// Sets the value 'v' for key 'k'.
	// k就是页号，v就是Span*
	void set(Number k, void *v)
	{
		array_[k] = v;
	}
};

// Two-level radix tree
// 双层基数树，两次哈希
// 双层基数树整体占用的空间和单层相同：4*2^5*2^15=4MB（32位下）
// 但是如果是单层基数树需要将4MB空间一次开好，而双层只需要先开好第一层，其余的空间用多少开多少
// 对于32位：4*2^5 + n*4*2^15；对于64位：4*2^5 + n*4*2^47(第二层的空间还是太大了)
template <int BITS>
class TCMalloc_PageMap2
{
private:
	// Put 32 entries in the root and (2^BITS)/32 entries in each leaf.
	// 取页号的前5位作为第一层索引
	static const int ROOT_BITS = 5;
	static const int ROOT_LENGTH = 1 << ROOT_BITS;
	// 剩下的位作为第二层索引
	static const int LEAF_BITS = BITS - ROOT_BITS;
	static const int LEAF_LENGTH = 1 << LEAF_BITS;

	// Leaf node
	struct Leaf
	{
		void *values[LEAF_LENGTH];
	};

	// 第一层直接写成静态数组
	Leaf *root_[ROOT_LENGTH]; // Pointers to 32 child nodes
							  // void *(*allocator_)(size_t); // Memory allocator

public:
	typedef uintptr_t Number;

	// explicit TCMalloc_PageMap2(void* (*allocator)(size_t)) {
	explicit TCMalloc_PageMap2()
	{
		// allocator_ = allocator;
		memset(root_, 0, sizeof(root_));

		PreallocateMoreMemory();
	}

	void *get(Number k) const
	{
		const Number i1 = k >> LEAF_BITS;		 // 取页号的高5位
		const Number i2 = k & (LEAF_LENGTH - 1); // 取页号的低15位
		if ((k >> BITS) > 0 || root_[i1] == NULL)
		{
			return NULL;
		}
		return root_[i1]->values[i2]; // 两次哈希
	}

	void set(Number k, void *v)
	{
		const Number i1 = k >> LEAF_BITS;
		const Number i2 = k & (LEAF_LENGTH - 1);
		assert(i1 < ROOT_LENGTH);
		root_[i1]->values[i2] = v;
	}

	// Ensure函数用于确保从start页号往后的n页，在基数树中的索引结构内存都已开好
	bool Ensure(Number start, size_t n)
	{
		for (Number key = start; key <= start + n - 1;)
		{
			const Number i1 = key >> LEAF_BITS;

			// Check for overflow
			if (i1 >= ROOT_LENGTH)
				return false;

			// Make 2nd level node if necessary
			if (root_[i1] == NULL)
			{
				// Leaf* leaf = reinterpret_cast<Leaf*>((*allocator_)(sizeof(Leaf)));
				// if (leaf == NULL) return false;
				static ObjectPool<Leaf> leafPool;
				Leaf *leaf = (Leaf *)leafPool.New();

				memset(leaf, 0, sizeof(*leaf));
				root_[i1] = leaf;
			}

			// Advance key past whatever is covered by this leaf node
			key = ((key >> LEAF_BITS) + 1) << LEAF_BITS;
		}
		return true;
	}

	// 直接把所有页的索引结构内存都开好
	void PreallocateMoreMemory()
	{
		// Allocate enough to keep track of all possible pages
		Ensure(0, 1 << BITS);
	}
};

// Three-level radix tree
// 三层基数树，三次哈希（适用于64位平台）
// 大小：对于64位：8*2^18 + m*8*2^18 + n*8*2^16（2MB + m*2MB + n*0.5MB）
template <int BITS>
class TCMalloc_PageMap3
{
private:
	// How many bits should we consume at each interior level
	// 前两层各占18位（64位下）
	static const int INTERIOR_BITS = (BITS + 2) / 3; // Round-up
	static const int INTERIOR_LENGTH = 1 << INTERIOR_BITS;

	// How many bits should we consume at leaf level
	// 第三层占16位
	static const int LEAF_BITS = BITS - 2 * INTERIOR_BITS;
	static const int LEAF_LENGTH = 1 << LEAF_BITS;

	// Interior node
	struct Node
	{
		Node *ptrs[INTERIOR_LENGTH];
	};

	// Leaf node
	struct Leaf
	{
		void *values[LEAF_LENGTH];
	};

	Node *root_; // Root of radix tree
	// void *(*allocator_)(size_t); // Memory allocator

	Node *NewNode()
	{
		// Node *result = reinterpret_cast<Node *>((*allocator_)(sizeof(Node)));
		size_t alignSize = AlineRule::_RoundUp(sizeof(Node), 1 << PAGE_SHIFT); // 按页对齐
		// cout << "PageMap Node alignSize: " << alignSize << endl; // debug
		Node *result = (Node *)SystemAlloc(alignSize >> PAGE_SHIFT);
		if (result != NULL)
		{
			memset(result, 0, sizeof(*result));
		}
		return result;
	}

public:
	typedef uintptr_t Number;

	// explicit TCMalloc_PageMap3(void *(*allocator)(size_t))
	explicit TCMalloc_PageMap3()
	{
		// allocator_ = allocator;
		root_ = NewNode(); // 先把第一层开好
	}

	void *get(Number k) const
	{
		const Number i1 = k >> (LEAF_BITS + INTERIOR_BITS);			// 取页号的高18位
		const Number i2 = (k >> LEAF_BITS) & (INTERIOR_LENGTH - 1); // 取页号的中间18位
		const Number i3 = k & (LEAF_LENGTH - 1);					// 取页号的低16位
		if ((k >> BITS) > 0 ||
			root_->ptrs[i1] == NULL || root_->ptrs[i1]->ptrs[i2] == NULL)
		{
			return NULL;
		}
		return reinterpret_cast<Leaf *>(root_->ptrs[i1]->ptrs[i2])->values[i3]; // 三次哈希
	}

	void set(Number k, void *v)
	{
		assert(k >> BITS == 0);
		const Number i1 = k >> (LEAF_BITS + INTERIOR_BITS);
		const Number i2 = (k >> LEAF_BITS) & (INTERIOR_LENGTH - 1);
		const Number i3 = k & (LEAF_LENGTH - 1);
		Ensure(k, 1); // 确保这一页的索引结构内存已经开好
		reinterpret_cast<Leaf *>(root_->ptrs[i1]->ptrs[i2])->values[i3] = v;
	}

	bool Ensure(Number start, size_t n)
	{
		for (Number key = start; key <= start + n - 1;)
		{
			const Number i1 = key >> (LEAF_BITS + INTERIOR_BITS);
			const Number i2 = (key >> LEAF_BITS) & (INTERIOR_LENGTH - 1);

			// Check for overflow
			if (i1 >= INTERIOR_LENGTH || i2 >= INTERIOR_LENGTH)
				return false;

			// Make 2nd level node if necessary
			if (root_->ptrs[i1] == NULL)
			{
				// cout << "Make 2nd level node" << endl; //degug
				Node *n = NewNode();
				if (n == NULL)
					return false;
				root_->ptrs[i1] = n;
			}

			// Make leaf node if necessary
			if (root_->ptrs[i1]->ptrs[i2] == NULL)
			{
				// cout << "SystemAlloc(sizeof(Leaf) >> PAGE_SHIFT)" << endl; //degug
				// Leaf *leaf = reinterpret_cast<Leaf *>((*allocator_)(sizeof(Leaf)));
				size_t alignSize = AlineRule::_RoundUp(sizeof(Leaf), 1 << PAGE_SHIFT); // 按页对齐
				// cout << "PageMap Leaf alignSize: " << alignSize << endl; //debug
				Leaf *leaf = (Leaf *)SystemAlloc(alignSize >> PAGE_SHIFT);
				if (leaf == NULL)
					return false;
				memset(leaf, 0, sizeof(*leaf));
				root_->ptrs[i1]->ptrs[i2] = reinterpret_cast<Node *>(leaf);
			}

			// Advance key past whatever is covered by this leaf node
			key = ((key >> LEAF_BITS) + 1) << LEAF_BITS;
		}
		return true;
	}

	void PreallocateMoreMemory()
	{
	}
};